// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/testing/exec/util/operator_test_base.h>
#include <pollux/parse/query_planner.h>

namespace kumo::pollux::exec::test {

class SqlTest : public OperatorTestBase {
 protected:
  void TearDown() override {
    planner_.reset();
    OperatorTestBase::TearDown();
  }

  void assertSql(const std::string& sql, const std::string& duckSql = "") {
    auto plan = planner_->plan(sql);
    AssertQueryBuilder(plan, duckDbQueryRunner_)
        .assertResults(duckSql.empty() ? sql : duckSql);
  }

  std::unique_ptr<core::DuckDbQueryPlanner> planner_{
      std::make_unique<core::DuckDbQueryPlanner>(pool())};
};

TEST_F(SqlTest, values) {
  assertSql("SELECT x, x + 5 FROM UNNEST([1, 2, 3]) as t(x)");
  assertSql("SELECT avg(x), count(*) FROM UNNEST([1, 2, 3]) as t(x)");
  assertSql("SELECT x % 5, avg(x) FROM UNNEST([1, 2, 3]) as t(x) GROUP BY 1");
  assertSql("SELECT avg(x * 4) FROM UNNEST([1, 2, 3]) as t(x)");
  assertSql(
      "SELECT x / 5, avg(x * 4) FROM UNNEST([1, 2, 3]) as t(x) GROUP BY 1");
}

TEST_F(SqlTest, customScalarFunctions) {
  planner_->registerScalarFunction(
      "array_join", {ARRAY(BIGINT()), VARCHAR()}, VARCHAR());

  assertSql("SELECT array_join([1, 2, 3], '-')", "SELECT '1-2-3'");
}

TEST_F(SqlTest, customAggregateFunctions) {
  // We need an aggregate that DuckDB does not support. 'every' fits the need.
  // 'every' is an alias for bool_and().
  planner_->registerAggregateFunction("every", {BOOLEAN()}, BOOLEAN());

  assertSql(
      "SELECT every(x) FROM UNNEST([true, false, true]) as t(x)",
      "SELECT false");
  assertSql(
      "SELECT x, every(x) FROM UNNEST([true, false, true, false]) as t(x) GROUP BY 1",
      "VALUES (true, true), (false, false)");
}

TEST_F(SqlTest, tableScan) {
  std::unordered_map<std::string, std::vector<RowVectorPtr>> data = {
      {"t",
       {make_row_vector(
           {"a", "b", "c"},
           {
               make_flat_vector<int64_t>({1, 2, 3, 1, 2, 3}),
               make_flat_vector<int32_t>({1, 2, 3, 4, 5, 6}),
               make_flat_vector<int16_t>({3, 6, 8, 2, 9, 10}),
           })}},
      {"u",
       {make_row_vector(
           {"a", "b"},
           {
               make_flat_vector<int64_t>({1, 2, 3, 4, 5, 1}),
               make_flat_vector<double>({1.1, 2.2, 3.3, 4.4, 5.5, 1.2}),
           })}},
  };

  createDuckDbTable("t", data.at("t"));
  createDuckDbTable("u", data.at("u"));

  planner_->registerTable("t", data.at("t"));
  planner_->registerTable("u", data.at("u"));

  assertSql("SELECT a, avg(b) FROM t WHERE c > 5 GROUP BY 1");
  assertSql("SELECT * FROM t, u WHERE t.a = u.a");
  assertSql("SELECT t.a, t.b, t.c, u.b FROM t, u WHERE t.a = u.a");
  assertSql("SELECT t.a, t.b, t.c, u.b FROM t left join u on t.a = u.a");
  assertSql(
      "SELECT t.a, t.b, t.c FROM t WHERE EXISTS (SELECT 1 FROM u WHERE t.a = u.a)");
  assertSql("SELECT t.a, t.b, t.c FROM t WHERE a < (SELECT max(u.a) FROM u)");
}

} // namespace kumo::pollux::exec::test
